import threading
import queue
import os
import time
import requests
from Crypto.Cipher import AES

class DownloadWorker(threading.Thread):
    def __init__(self, task_queue, db, task_id, task_uid, save_path, log_callback, key=None):
        threading.Thread.__init__(self)
        self.task_queue = task_queue
        self.db = db
        self.task_id = task_id
        self.task_uid = task_uid
        self.save_path = save_path
        self.log_callback = log_callback
        self.key = key
        self.running = True
        
    def run(self):
        """线程主函数"""
        while self.running:
            segment = self.get_next_segment()
            if segment is None:
                break
                
            self.download_segment(segment)
            
        self.log_callback(f"工作线程 {threading.current_thread().name} 退出")
    
    def get_next_segment(self):
        """从队列获取下一个片段"""
        try:
            return self.task_queue.get(timeout=1)
        except queue.Empty:
            return None
    
    def download_segment(self, segment):
        """下载单个TS片段"""
        ts_path = os.path.join(self.save_path, f"{self.task_uid}_{segment['order_num']}.ts")
        
        try:
            # 下载TS片段
            headers = {"Referer": segment.get('refer', '')}
            response = requests.get(segment['link'], headers=headers, timeout=30)
            
            data = response.content
            
            # 解密处理
            if self.key:
                iv = segment['order_num'].to_bytes(16, 'big')  # 使用序号作为IV
                cipher = AES.new(self.key, AES.MODE_CBC, iv=iv)
                data = cipher.decrypt(data)
            
            # 保存文件
            with open(ts_path, 'wb') as f:
                f.write(data)
            
            # 更新数据库状态
            self.db.update_ts_state(segment['id'], 1)
            self.log_callback(f"片段 {segment['order_num']} 下载完成")
            
        except Exception as e:
            self.log_callback(f"片段下载失败: {str(e)}")
            # 重新加入队列
            self.task_queue.put(segment)
        finally:
            self.task_queue.task_done()

class DownloadManager:
    def __init__(self, db, max_workers=5):
        self.db = db
        self.max_workers = max_workers
        self.active_tasks = {}
        
    def start_download(self, task_id, task_uid, save_path, log_callback):
        """启动下载任务"""
        if task_id in self.active_tasks:
            log_callback(f"任务 {task_id} 已在下载中")
            return
            
        # 创建下载队列
        task_queue = queue.Queue()
        
        # 获取任务信息
        task_info = self.db.get_task(task_id)
        
        # 下载密钥
        key = None
        if task_info['keyUrl']:
            log_callback("下载密钥...")
            try:
                response = requests.get(task_info['keyUrl'], headers={"Referer": task_info['refer']})
                key = response.content
                log_callback(f"密钥下载成功: {len(key)}字节")
            except Exception as e:
                log_callback(f"密钥下载失败: {str(e)}")
        
        # 创建保存目录
        save_dir = os.path.join(save_path, task_info['name'])
        os.makedirs(save_dir, exist_ok=True)
        
        # 添加所有待下载片段到队列
        segments = self.db.get_ts_segments(task_uid)
        for segment in segments:
            task_queue.put(segment)
        
        # 创建工作线程
        workers = []
        for i in range(self.max_workers):
            worker = DownloadWorker(
                task_queue, 
                self.db, 
                task_id,
                task_uid,
                save_dir,
                log_callback,
                key
            )
            worker.daemon = True
            worker.start()
            workers.append(worker)
        
        # 保存任务状态
        self.active_tasks[task_id] = {
            'queue': task_queue,
            'workers': workers,
            'task_info': task_info
        }
        
        log_callback(f"开始下载任务 {task_id}，启动 {self.max_workers} 个工作线程")
        
        # 启动监控线程
        monitor_thread = threading.Thread(
            target=self.monitor_download,
            args=(task_id, task_uid, save_dir, log_callback)
        )
        monitor_thread.daemon = True
        monitor_thread.start()
    
    def pause_download(self, task_id):
        """暂停下载任务"""
        if task_id in self.active_tasks:
            workers = self.active_tasks[task_id]['workers']
            for worker in workers:
                worker.running = False
            
            del self.active_tasks[task_id]
    
    def monitor_download(self, task_id, task_uid, save_dir, log_callback):
        """监控下载任务完成状态"""
        while task_id in self.active_tasks:
            # 检查队列是否为空且所有工作已完成
            if (self.active_tasks[task_id]['queue'].empty() and 
                all(not worker.is_alive() for worker in self.active_tasks[task_id]['workers'])):
                
                # 检查数据库确认所有片段已完成
                if self.db.check_task_completed(task_uid):
                    self.db.mark_task_completed(task_id)
                    log_callback(f"任务 {task_id} 下载完成")
                    
                    # 合并文件
                    self.merge_files(task_id, save_dir, log_callback)
                
                # 移除任务
                if task_id in self.active_tasks:
                    del self.active_tasks[task_id]
                break
                
            time.sleep(2)
    
    def merge_files(self, task_id, save_dir, log_callback):
        """合并所有TS片段为单个文件"""
        log_callback("开始合并文件...")
        task_info = self.db.get_task(task_id)
        
        # 获取所有已下载的TS片段路径
        ts_files = sorted(
            [f for f in os.listdir(save_dir) if f.startswith(task_info['uid']) and f.endswith('.ts')],
            key=lambda x: int(x.split('_')[-1].split('.')[0])
        )
        
        # 最终视频路径
        output_file = os.path.join(
            save_dir,
            f"{task_info['name']}.mp4"
        )
        
        # 合并文件
        try:
            with open(output_file, 'wb') as out_f:
                for ts_file in ts_files:
                    ts_path = os.path.join(save_dir, ts_file)
                    with open(ts_path, 'rb') as in_f:
                        out_f.write(in_f.read())
                    # 清理临时文件
                    os.remove(ts_path)
            
            log_callback(f"合并成功: {output_file}")
        except Exception as e:
            log_callback(f"合并失败: {str(e)}")
    
    def update_ui_status(self):
        """更新UI状态（在主线程中调用）"""
        # 在实际应用中，这里应调用UI更新方法
        pass
    
    def cleanup_task_files(self, task_id):
        """清理任务文件"""
        # 如果任务正在下载，先停止
        self.pause_download(task_id)
        
        # 清理文件
        # 注意: 文件清理现在由主程序处理
        pass
        
    def shutdown(self):
        """关闭下载管理器，停止所有活动任务"""
        # 停止所有活动的下载任务
        for task_id in list(self.active_tasks.keys()):
            self.pause_download(task_id)
        
        # 确保所有工作线程已停止
        for task_data in self.active_tasks.values():
            for worker in task_data['workers']:
                if worker.is_alive():
                    worker.running = False
                    worker.join(timeout=1.0)
        
        # 清空活动任务列表
        self.active_tasks.clear()